"""
LangChain 的回调系统，其核心思想是：在 LLM 应用执行的各个关键节点（或称为“事件”），向用户提供一个“钩子”（Hook），让用户能够插入自定义的逻辑。
这就像是在程序的流水线上安装了一系列的监控探头，当流水线经过某个探头时，就会触发一个信号，执行你预设的代码。
LangChain 定义了一系列标准化的回调事件，覆盖了执行过程的方方面面。以下是一些最常见和重要的事件：
    on_chain_start / on_chain_end： 当一个链（Chain）开始/结束执行时触发。
    on_llm_start / on_llm_end：当对大语言模型（LLM）的调用开始/结束时触发。
    on_llm_end 会返回 LLM 生成的 Token 数量、耗时等信息。
    on_chat_model_start：类似于 on_llm_start，但专门用于聊天模型。
    on_tool_start / on_tool_end：当代理（Agent）开始/结束执行一个工具（Tool）时触发。
    on_agent_action：当代理决定要采取哪个行动（即使用哪个工具）时触发。
    on_retriever_start / on_retriever_end：当检索器（Retriever）开始/结束检索时触发
"""

import os

from dotenv import load_dotenv
from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.callbacks import StreamingStdOutCallbackHandler
from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

# 加载环境变量
load_dotenv()

# 多个专用回调处理器
class LoggingCallbackHandler(BaseCallbackHandler):
    def on_chain_start(self, serialized, inputs, **kwargs):
        print(f"📝 [日志] 开始处理: {list(inputs.keys())}")

    def on_chain_end(self, outputs, **kwargs):
        print(f"📝 [日志] 处理完成")


class TimingCallbackHandler(BaseCallbackHandler):
    def __init__(self):
        self.start_time = None

    def on_llm_start(self, serialized, prompts, **kwargs):
        self.start_time = os.times().elapsed

    def on_llm_end(self, response, **kwargs):
        end_time = os.times().elapsed
        duration = end_time - self.start_time
        print(f"⏱️ [计时] LLM 耗时: {duration:.2f}秒")


class TokenCounterCallbackHandler(BaseCallbackHandler):
    def on_llm_end(self, response, **kwargs):
        usage = response.llm_output.get('token_usage', {})
        print(f"🔢 [令牌] 使用: {usage.get('prompt_tokens', 0)} 输入 + {usage.get('completion_tokens', 0)} 输出 = {usage.get('total_tokens', 0)} 总计")


# 创建链 之定义chatModel （自定义concept参数和example_count）
template = PromptTemplate.from_template("请解释{concept}，并给出{example_count}个例子")
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen2.5-VL-7B-Instruct", streaming=True, temperature=0.7)
chain = LLMChain(llm=llm, prompt=template)

print("=== 请求: 完整监控 ===")
response1 = chain.invoke(
    {"concept": "神经网络", "example_count": 2},
    config={
        "callbacks": [
            LoggingCallbackHandler(),
            TimingCallbackHandler(),
            TokenCounterCallbackHandler(),
            StreamingStdOutCallbackHandler()  # 流式输出
        ]
    }
)

